package com.yjz.aiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * @author YJZ
 * @Description 自定义 VectorStore 的整合配置类
 */
@Configuration
public class PGVectorVectorStoreConfig {

    @Resource
    private LoveAppDocumentLoader loveAppDocumentLoader;

    // EmbeddingModel dashscopeEmbeddingModel 入参异地昂要是精确的阿里云的 EmbeddingModel
    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate , EmbeddingModel dashscopeEmbeddingModel) {
        PgVectorStore vetorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)
                .distanceType(COSINE_DISTANCE)
                .indexType(HNSW)
                .initializeSchema(true)
                .schemaName("public")
                .vectorTableName("vector_store")
                .maxDocumentBatchSize(10000)
                .build();

        List<Document> documents = loveAppDocumentLoader.loadMarkdowns();
        vetorStore.add(documents);
        return vetorStore;
    }

}
